/*
 * H265RtpPackager.cpp
 *
 *  Created on: 2016年4月23日
 *      Author: terry
 */

#include "H265RtpPackager.h"
#include "H264NaluParser.h"
#include "RtpHeader.h"


namespace av
{

H265RtpPackager::H265RtpPackager():
		m_lastTime()
{
}

H265RtpPackager::~H265RtpPackager()
{
}


void H265RtpPackager::slice(const MediaPacket& pkt, int maxSize, RtpPackagerSink* pSink)
{
    if (!pSink)
    {
        return;
    }

    uint8_t* data = pkt.data;
    int size = pkt.size;

    int count = 0;
    if (H264NaluParser::startWithH264Code(data, size))
    {
        size_t begin = 0;
        bool found = true;

        while (found)
        {
            size_t pos = H264NaluParser::findH264StartCode(data, size, begin + 3);
            if (pos != (size_t)-1)
            {
                sliceNaluPtr(data + begin, data + pos, (uint32_t)pkt.pts, false, maxSize, pSink);
                begin = pos;
                count ++;
            }
            else
            {
                break;
            }
        }

        sliceNaluPtr(data + begin, data + size, (uint32_t)pkt.pts, true, maxSize, pSink);
        count ++;
    }
    else
    {
        sliceNalu(data, size, (uint32_t)pkt.pts, true, maxSize, pSink);
        count ++;
    }
}

void H265RtpPackager::sliceNaluPtr(uint8_t* data, uint8_t* end, uint32_t timestamp, bool mark,
        int maxSize, RtpPackagerSink* pSink)
{
    if (end <= data)
    {
        return;
    }

    size_t size = (end - data);
    NaluPacket nalu;
    if (H264NaluParser::parseNalu(data, size, nalu))
    {
        sliceNalu(data + nalu.prefix, size - nalu.prefix, timestamp, mark, maxSize, pSink);
    }
    else
    {
        sliceNalu(data, size, timestamp, mark, maxSize, pSink);
    }
}

void H265RtpPackager::sliceNalu(uint8_t* data, int size, uint32_t timestamp, bool mark,
               int maxSize, RtpPackagerSink* pSink)
{
	maxSize -= 2;

    if (size <= maxSize)
    {
        RtpPacket pkt;
        pkt.ts = timestamp;
        pkt.mark = mark;
        pkt.data = data;
        pkt.size = size;
		pkt.pt = m_payload;

        pSink->onSlicePacket(pkt);
        return;
    }

    uint8_t naluType = (data[0] & 0x7E) >> 1;

    uint8_t hdr[2];
    hdr[0] = (data[0] & 0x81) | (49<<1);
    hdr[1] = data[1];

    data += 2;
    size -= 2;

    RtpPacket pkt;
    pkt.ts = timestamp;
    pkt.mark = false;
	pkt.pt = m_payload;

    bool start = true;
    bool end = false;
    while (size > maxSize)
    {
        m_sliceBuffer.clear();
        m_sliceBuffer.write(hdr[0]);
        m_sliceBuffer.write(hdr[1]);
        m_sliceBuffer.write(naluType);
        m_sliceBuffer.write(data, maxSize);

        pkt.data = m_sliceBuffer.data();
        pkt.size = m_sliceBuffer.size();

        FU_HEADER* pFuHeader = (FU_HEADER*)(pkt.data + 2);
        pFuHeader->TYPE = naluType & 0x1F;
        pFuHeader->S   = start;
        pFuHeader->E   = end;
        pFuHeader->R   = 0;

        pSink->onSlicePacket(pkt);

        data += maxSize;
        size -= maxSize;

        start = false;
    }

    // last packet
    pkt.mark = mark;

    end = true;

    m_sliceBuffer.clear();
    m_sliceBuffer.write(hdr[0]);
    m_sliceBuffer.write(hdr[1]);
    m_sliceBuffer.write(naluType);
    m_sliceBuffer.write(data, size);

    pkt.data = m_sliceBuffer.data();
    pkt.size = m_sliceBuffer.size();

    FU_HEADER* pFuHeader = (FU_HEADER*)(pkt.data + 2);
    pFuHeader->TYPE = naluType & 0x1F;
    pFuHeader->S            = start;
    pFuHeader->E            = end;
    pFuHeader->R            = 0;

    pSink->onSlicePacket(pkt);
}

bool H265RtpPackager::join(const RtpPacket& pktIn, MediaPacket& pktOut)
{
    static unsigned char s_startCode[] = { 0, 0, 0, 1};

    uint8_t* data = pktIn.data;
    int length = pktIn.size;

    bool isFU = ((data[0] >> 1) & 49) == 49;
    uint8_t naluType = (data[0] & 0x7E) >> 1;
    bool fullPacket = pktIn.mark;
    bool keyFrame = false;

    if (isFU)
    {
    	naluType = data[2] & 0x1F;

    	uint8_t hdr[2] = {0, data[1]};
    	hdr[0] = (data[0] & 0x81) | ((naluType << 1) & 0x7E);

        unsigned char startBit = data[2]&0x80;
        unsigned char endBit = data[2]&0x40;

        if(startBit)
        {
            m_joinBuffer.clear();
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
            m_joinBuffer.write(hdr, sizeof(hdr));
        }
        else if (m_joinBuffer.empty())
        {
            m_joinBuffer.clear();
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
            m_joinBuffer.write(hdr, sizeof(hdr));
        }

        data += 3;
        length -= 3;

        m_joinBuffer.write(data, length);

        if (endBit)
        {
            fullPacket = true;
        }
    }
    else
    {
        if (m_joinBuffer.empty())
        {
            m_joinBuffer.write(s_startCode, sizeof(s_startCode));
        }
        else
        {
            //size_t len = m_joinBuffer.length();
        }

        m_joinBuffer.write(data, length);
        fullPacket = true;
    }

    if (fullPacket)
    {
        pktOut.data = m_joinBuffer.data();
        pktOut.size = m_joinBuffer.length();
        pktOut.type = MEDIA_TYPE_VIDEO;
        pktOut.pts = pktIn.ts;
        pktOut.duration = 0;


        if (keyFrame)
        {
            pktOut.flags |= MEDIA_FLAG_KEY;
        }

        m_joinBuffer.clear();
        m_lastTime = pktIn.ts;
    }

    return fullPacket;
}


void H265RtpPackager::reset()
{
    m_sliceBuffer.clear();

    m_joinBuffer.clear();
    m_lastTime = 0;
}

bool H265RtpPackager::getFormat(MediaFormat& fmt)
{
	fmt.m_codec = MEDIA_CODEC_HEVC;
	return true;
}


} /* namespace av */
